import os
import sys
import importlib.util
from google.adk.agents import Agent, LlmAgent
from google.adk.tools import FunctionTool
from google.adk.models.lite_llm import LiteLlm
from typing import Dict, Any

# 定义动态导入函数
def import_module_from_path(module_name, file_path):
    """动态导入模块"""
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module

# 确定当前目录
current_dir = os.path.dirname(os.path.abspath(__file__))

try:
    # 尝试普通导入
    from knowledge_tool import knowledge_retrieval_tool, initialize_knowledge_base_with_examples
except ImportError:
    # 如果失败，尝试动态导入
    knowledge_tool_module = import_module_from_path(
        "knowledge_tool", os.path.join(current_dir, "knowledge_tool.py")
    )
    knowledge_retrieval_tool = knowledge_tool_module.knowledge_retrieval_tool
    initialize_knowledge_base_with_examples = knowledge_tool_module.initialize_knowledge_base_with_examples


# 将Python函数封装为ADK工具
knowledge_tool_instance = FunctionTool(knowledge_retrieval_tool)

# API密钥设置
DEFAULT_DASHSCOPE_API_KEY = "sk-f227634bb56*************************232"  # 请替换为您的实际密钥
DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY", DEFAULT_DASHSCOPE_API_KEY)


# 企业知识库问答代理
enterprise_qna_agent = LlmAgent(
    name="EnterpriseKnowledgeAssistant",
    model=LiteLlm(
                model="openai/qwen-turbo",
                api_key=DASHSCOPE_API_KEY,
                api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    ),
    tools=[knowledge_tool_instance],
    instruction="""
    你是一个企业内部知识问答助手。你的任务是根据用户的问题，利用 knowledge_retrieval_tool 从公司知识库中查找相关信息，并基于这些信息清晰、准确地回答用户。

    请遵循以下步骤：
    1. 分析用户的问题，理解其意图。
    2. 如果问题适合从知识库中查找答案，请立即调用 knowledge_retrieval_tool，并将用户的完整问题作为查询参数。
    3. 工具会返回相关的文档片段。仔细阅读这些片段。
    4. 严格根据检索到的文档片段内容来组织你的回答。不要添加外部知识或个人猜测。
    5. 如果检索到的文档片段包含答案，请清晰地呈现。你可以引用关键信息或总结要点。
    6. 如果工具返回 'status: "not_found"' 或检索到的文档与问题不相关，请明确告知用户你在公司知识库中没有找到直接答案，不要尝试编造。
    7. 保持回答专业、简洁、有帮助。

    示例：
    用户："我们公司的年假政策是什么？"
    你的行动：[调用 knowledge_retrieval_tool(query="我们公司的年假政策是什么？")]
    (假设工具返回了相关政策片段)
    你的回答：[基于返回的片段进行回答，例如："根据公司政策文档，员工每年享有XX天年假..."]

    用户："如何配置新的开发环境？"
    你的行动：[调用 knowledge_retrieval_tool(query="如何配置新的开发环境？")]
    (假设工具返回了相关设置指南)
    你的回答：[基于返回的片段进行回答，例如："配置新的开发环境，请参照以下步骤..."]
    """,
    output_key="qna_response"
)


# 为兼容 ADK 的要求，将智能体公开为 root_agent
root_agent = enterprise_qna_agent


# 在导入模块时初始化知识库
try:
    initialize_knowledge_base_with_examples()
except Exception as e:
    print(f"初始化知识库时出错: {e}")
    print("请确保已安装必要的依赖: pip install pymilvus 'pymilvus[model]'")